package com.ruoyi.websocket.config;

import com.ruoyi.websocket.domain.AutoReply;
import org.elasticsearch.index.query.QueryBuilders;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArraySet;

public class MyWebSocketHandler extends TextWebSocketHandler {

    // 存储所有连接的会话
    private static final CopyOnWriteArraySet<WebSocketSession> sessions = new CopyOnWriteArraySet<>();
    // 存储用户会话和身份信息的映射
    private static final Map<WebSocketSession, Boolean> sessionIdentityMap = new HashMap<>();

    @Autowired
    private ElasticsearchOperations elasticsearchOperations;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        System.out.println("WebSocket 连接已建立，会话 ID: " + session.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        try {
            // 处理身份验证消息
            if (payload.startsWith("{\"type\": \"identity\"")) {
                handleIdentityMessage(session, payload);
                return;
            }

            // 处理普通消息
            handleNormalMessage(session, payload);
        } catch (Exception e) {
            System.err.println("处理消息时出错: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private void handleIdentityMessage(WebSocketSession session, String payload) {
        try {
            // 解析身份信息
            // 这里可以使用 JSON 解析库，如 Jackson 或 Gson
            boolean isCustomerService = payload.contains("\"isCustomerService\": true");
            sessionIdentityMap.put(session, isCustomerService);
            System.out.println("身份验证成功，会话 ID: " + session.getId() + ", 是否为客服: " + isCustomerService);
        } catch (Exception e) {
            System.err.println("身份验证消息解析失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private void handleNormalMessage(WebSocketSession session, String payload) throws IOException {
        boolean isCustomerService = sessionIdentityMap.getOrDefault(session, false);
        if (isCustomerService) {
            // 客服发送的消息，转发给用户
            for (Map.Entry<WebSocketSession, Boolean> entry : sessionIdentityMap.entrySet()) {
                WebSocketSession targetSession = entry.getKey();
                boolean targetIsCustomerService = entry.getValue();
                if (!targetIsCustomerService && targetSession.isOpen()) {
                    targetSession.sendMessage(new TextMessage(payload));
                }
            }
        } else {
            // 用户发送的消息，先尝试自动回复
            String answer = searchAnswerInES(payload);
            if (answer != null) {
                // 如果检索到答案，发送自动回复消息给用户，并在前面加上 [自动回复]
                answer = "[自动回复] " + answer;
                String autoReplyMessage = "{\"type\": \"auto_reply\", \"text\": \"" + answer + "\"}";
                session.sendMessage(new TextMessage(autoReplyMessage));
            } else {
                // 未检索到答案，转发给客服
                for (Map.Entry<WebSocketSession, Boolean> entry : sessionIdentityMap.entrySet()) {
                    WebSocketSession targetSession = entry.getKey();
                    boolean targetIsCustomerService = entry.getValue();
                    if (targetIsCustomerService && targetSession.isOpen()) {
                        targetSession.sendMessage(new TextMessage(payload));
                    }
                }
            }
        }
    }

    private String searchAnswerInES(String question) {
        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder()
                .withQuery(QueryBuilders.matchQuery("question", question))
                .build();
        SearchHits<AutoReply> searchHits = elasticsearchOperations.search(searchQuery, AutoReply.class);
        if (!searchHits.isEmpty()) {
            SearchHit<AutoReply> searchHit = searchHits.getSearchHit(0);
            AutoReply autoReply = searchHit.getContent();
            return autoReply.getAnswer();
        }
        return null;
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        sessions.remove(session);
        sessionIdentityMap.remove(session);
        System.out.println("WebSocket 连接已关闭，会话 ID: " + session.getId());
    }
}